package com.jiyun.ai.controller;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.Constants;
import io.reactivex.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Arrays;

@RestController
@RequestMapping("/chat")
public class ChatController {
    private static final Logger logger = LoggerFactory.getLogger(ChatController.class);

    @Value("${ai.api_key}")
    private String apiKey;

    @RequestMapping(value = "/ask", produces = MediaType.TEXT_EVENT_STREAM_VALUE )
    public SseEmitter callMany(@RequestParam(value = "question", required = false) String message) throws NoApiKeyException, InputRequiredException {

        try {
            // 设置API密钥
            Constants.apiKey = apiKey;
            Generation gen = new Generation();
            Message userMsg = Message.builder().role(Role.USER.getValue()).content(message).build();

            Flowable<GenerationResult> serverSentEventFlux = streamCallWithMessage(gen, userMsg);

            // 创建 SseEmitter
            SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
            // 订阅 Flowable 并发送数据到 SseEmitter
            serverSentEventFlux.subscribe(result -> sendSseData(emitter, result.getOutput().getChoices().get(0).getMessage().getContent()));

            // 当完成时关闭 SseEmitter
            emitter.onCompletion(() -> {
                System.out.println("SSE 连接已完成，正在关闭.");
            });

            return emitter;

        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            logger.error("发生异常: {}", e.getMessage());
        }
        return null;
    }

    // 发送数据到 SseEmitter
    private void sendSseData(SseEmitter emitter, String data) {
        try {
            // 发送数据到 SseEmitter
            emitter.send(data.toString());
        } catch (Exception e) {
            // 处理错误情况
            System.err.println("发送数据失败: " + e.getMessage());
            emitter.completeWithError(e);
        }
    }

    // 流式调用
    public  Flowable<GenerationResult> streamCallWithMessage(Generation gen, Message userMsg)
            throws NoApiKeyException, ApiException, InputRequiredException {
        // 根据用户消息构建生成参数
        GenerationParam param = buildGenerationParam(userMsg);
        // 使用构建的参数启动生成调用，并返回生成结果的Flowable对象
        Flowable<GenerationResult> result = gen.streamCall(param);
        // 从结果创建一个Flux流
        return result;
    }

    // 构建 GenerationParam
    private GenerationParam buildGenerationParam(Message userMsg) {
        return GenerationParam.builder()
                .model("qwen-turbo")
                .messages(Arrays.asList(userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .incrementalOutput(true)
                .build();
    }
}